package org.lemon.transport.mqtt.server.service;

import io.netty.handler.ssl.SslHandler;
import io.netty.util.internal.StringUtil;
import lombok.extern.slf4j.Slf4j;
import org.lemon.transport.context.config.ssl.SslCredentials;
import org.lemon.transport.mqtt.config.SslCredentialsConfig;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;

import javax.net.ssl.*;

@Slf4j
@Component
public class MqttSslHandlerProvider {

    @Value("${transport.mqtt.ssl.protocol}")
    private String sslProtocol;

    @Autowired
    private SslCredentialsConfig sslCredentialsConfig;

    private SSLContext sslContext;

    public SslHandler getSslHandler() {

        if (sslContext == null) {
            this.sslContext = createSslContext();
        }
        SSLEngine sslEngine = sslContext.createSSLEngine();
        sslEngine.setUseClientMode(false);
        sslEngine.setNeedClientAuth(false);
        sslEngine.setWantClientAuth(true);
        sslEngine.setEnabledProtocols(sslEngine.getSupportedProtocols());
        sslEngine.setEnabledCipherSuites(sslEngine.getSupportedCipherSuites());
        sslEngine.setEnableSessionCreation(true);
        return new SslHandler(sslEngine);
    }

    private SSLContext createSslContext() {

        try {
            SslCredentials sslCredentials = this.sslCredentialsConfig.getCredentials();
            TrustManagerFactory trustManagerFactory = sslCredentials.createTrustManagerFactory();
            KeyManagerFactory keyManagerFactory = sslCredentials.createKeyManagerFactory();

            KeyManager[] keyManagers = keyManagerFactory.getKeyManagers();
            TrustManager x509TrustManager = getX509TrustManager(trustManagerFactory);
            TrustManager[] trustManagers = {x509TrustManager};
            if (StringUtil.isNullOrEmpty(sslProtocol)) {
                sslProtocol = "TLS";
            }
            SSLContext sslContext = SSLContext.getInstance(sslProtocol);
            sslContext.init(keyManagers, trustManagers, null);
            return sslContext;
        } catch (Exception e) {
            log.error("Unable to set up SSL context. Reason: " + e.getMessage(), e);
            throw new RuntimeException("Failed to get SSL context", e);
        }
    }

    private TrustManager getX509TrustManager(TrustManagerFactory tmf) throws Exception {

        X509TrustManager x509Tm = null;
        for (TrustManager tm : tmf.getTrustManagers()) {
            if (tm instanceof X509TrustManager) {
                x509Tm = (X509TrustManager) tm;
                break;
            }
        }
        return x509Tm;
    }
}
